package com.weixin.project.Config.Filter;

import org.springframework.http.HttpStatus;
import org.springframework.web.filter.OncePerRequestFilter;

import javax.servlet.FilterChain;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

import java.time.LocalDateTime;
import java.time.temporal.ChronoUnit;

public class RateLimitFilter extends OncePerRequestFilter {

    private static class RequestInfo {
        private AtomicInteger count;
        private LocalDateTime lastRequestTime;

        RequestInfo(int count, LocalDateTime lastRequestTime) {
            this.count = new AtomicInteger(count);
            this.lastRequestTime = lastRequestTime;
        }
    }

    private final ConcurrentHashMap<String, RequestInfo> requestCounts = new ConcurrentHashMap<>();
    private final int maxRequests;
    private final long timeWindowMillis;

    public RateLimitFilter(int maxRequests, long timeWindowMillis) {
        this.maxRequests = maxRequests;
        this.timeWindowMillis = timeWindowMillis;
    }

    @Override
    protected void doFilterInternal(HttpServletRequest request, HttpServletResponse response, FilterChain filterChain)
            throws ServletException, IOException {
        if (shouldFilter(request)){
//            String userId = request.getHeader("user_id");
            String userId = request.getHeader("token");
            if (userId == null) {
                response.setStatus(HttpStatus.BAD_REQUEST.value());
                return;
            }

            RequestInfo requestInfo = requestCounts.computeIfAbsent(userId, k -> new RequestInfo(0, LocalDateTime.now()));

            if (ChronoUnit.MILLIS.between(requestInfo.lastRequestTime, LocalDateTime.now()) > timeWindowMillis) {
                requestInfo.count.set(0);
                requestInfo.lastRequestTime = LocalDateTime.now();
            }

            if (requestInfo.count.incrementAndGet() > maxRequests) {
                response.setStatus(HttpStatus.TOO_MANY_REQUESTS.value());
                return;
            }
        }
        filterChain.doFilter(request, response);
    }

    // 判断请求是否需要进行限流处理
    private boolean shouldFilter(HttpServletRequest request) {
        // 应该只有访问到题目这种付费资源时，才需要进行限流处理
        return request.getRequestURI().startsWith("/xxxx");
    }
}
